module Zeno.Term where

import Prelude ()
import StdImports
import qualified Data.Set as Set

data Term a
  = Var !a
  | App !(Term a) !(Term a)
  | Bot
  deriving ( Eq, Ord )
  
instance NFData a => NFData (Term a) where
  rnf Bot = ()
  rnf (Var a) = rnf a
  rnf (App t1 t2) = rnf t1 `seq` rnf t2

instance Functor Term where
  fmap f (Var x) = Var (f x)
  fmap f (App t1 t2) = App (fmap f t1) (fmap f t2)
  
instance Foldable Term where
  foldMap f (Var x) = f x
  foldMap f (App t1 t2) = foldMap f t1 `mappend` foldMap f t2
  
instance Traversable Term where
  sequenceA (Var x) = Var <$> x
  sequenceA (App t1 t2) = App <$> sequenceA t1 <*> sequenceA t2
  
class TermTraversable t where
  mapTermsM :: Monad m => (Term a -> m (Term a)) -> t a -> m (t a)
  
instance TermTraversable Term where
  mapTermsM f (App x y) =
    f =<< App `liftM` mapTermsM f x `ap` mapTermsM f y
  mapTermsM f t = f t

varTerm :: Term a -> Bool
varTerm (Var {}) = True
varTerm _ = False

fromVar :: Term a -> a
fromVar (Var a) = a

appTerm :: Term a -> Bool
appTerm (App {}) = True
appTerm _ = False
  
mapTerms :: TermTraversable t => (Term a -> Term a) -> t a -> t a
mapTerms = mapM_to_fmap mapTermsM

termList :: TermTraversable t => t a -> [Term a]
termList = execWriter . mapTermsM (\t -> tell [t] >> return t)
  
subTermList :: Eq a => Term a -> [Term a]
subTermList term = filter (/= term) (termList term)

containsTerm :: (TermTraversable t, Eq a) => t a -> Term a -> Bool
containsTerm t = flip elem (termList t)

containsSubTerm :: Eq a => Term a -> Term a -> Bool
containsSubTerm t = flip elem (subTermList t)

instance Show a => Show (Term a) where
  show (Var a) = show a
  show (App x (Var y)) = show x ++ " " ++ show y
  show (App x y) = show x ++ " (" ++ show y ++ ")"
  
replaceTerms :: (TermTraversable t, Eq a) => Term a -> Term a -> t a -> t a
replaceTerms = genericReplace mapTerms

termFunction :: Term a -> a
termFunction (Var a) = a
termFunction (App t1 t2) = termFunction t1

flattenTerm :: Term a -> [Term a]
flattenTerm (App f a) = flattenTerm f ++ [a]
flattenTerm t = [t]

unflattenTerm :: [Term a] -> Term a
unflattenTerm = foldl1 App

splitTerms :: (Term a -> [Term a]) -> Term a -> [Term a]
splitTerms f t@(App t1 t2) =
  let t1s = splitTerms f t1
      t2s = splitTerms f t2
  in f t ++ (map (App t1) t2s) ++ (map (flip App t2) t1s)
splitTerms f t = f t
  
splitTermsM :: Monad m => (Term a -> m [Term a]) -> Term a -> m [Term a]
splitTermsM f t@(App t1 t2) = do
  t1s <- splitTermsM f t1
  t2s <- splitTermsM f t2
  ft <- f t
  return $ ft ++ (map (App t1) t2s) ++ (map (flip App t2) t1s)
splitTermsM f t = f t

splitMapTerms :: Functor f => 
  (Term a -> [f (Term a)]) -> Term a -> [f (Term a)]
splitMapTerms f t@(App t1 t2) = 
  let t1s = splitMapTerms f t1
      t2s = splitMapTerms f t2
  in f t ++ (map (fmap (App t1)) t2s) ++ (map (fmap (flip App t2)) t1s) 
splitMapTerms f t = f t 

splitMapTermsM :: (Functor f, Monad m) =>
  (Term a -> m [f (Term a)]) -> Term a -> m [f (Term a)]
splitMapTermsM f t@(App t1 t2) = do
  t1s <- splitMapTermsM f t1
  t2s <- splitMapTermsM f t2
  t' <- f t
  return $ t' ++ (map (fmap (App t1)) t2s) 
    ++ (map (fmap (flip App t2)) t1s)
splitMapTermsM f t = f t
  
type Substitution a = Set (a, a)

applySubstitution :: (Functor f, Eq a) => 
  Substitution a -> f a -> f a
applySubstitution = flip $ foldl' $ \f (v1, v2) -> replace v1 v2 f

type TermSubstitution a = Substitution (Term a)

applyTermSubstitution :: (TermTraversable t, Eq a) => 
  TermSubstitution a -> t a -> t a
applyTermSubstitution = flip $ foldl' $ \tt (t1, t2) -> replaceTerms t1 t2 tt

normalise :: Ord a => TermSubstitution a -> TermSubstitution a
normalise sub = 
  let sub' = Set.map (\(x, y) -> 
        (applyTermSubstitution sub x, applyTermSubstitution sub y)) sub
  in if sub' /= sub then normalise sub' else sub'

unify :: Ord a => Term a -> Term a -> TermSubstitution a
unify t1 t2 =
  let (t1_func : t1_args) = flattenTerm t1
      (t2_func : t2_args) = flattenTerm t2
  in if (t1_func == t2_func) && (length t1_args == length t2_args)
       then mconcat (zipWith unify t1_args t2_args)
       else Set.singleton (t1, t2)
       
allUnifiers :: Ord a => Term a -> Term a -> [TermSubstitution a]
allUnifiers from =
  let unifiers to = tell [unify from to] >> return to
  in execWriter . (mapTermsM unifiers)
      
isOneToOne :: Eq a => Substitution a -> Bool
isOneToOne sub = 
  let (_, rs) = unzip $ Set.toList sub in
  isOneToMany sub && not (containsDuplicates rs)
  
isOneToMany :: Eq a => Substitution a -> Bool
isOneToMany sub =                       
  let (ls, _) = unzip $ Set.toList sub in 
  not (containsDuplicates ls) 

